

#!/bin/bash

# Usage: ./test.sh -c config.yaml
######################################################################################
# Step 0: Init
# Author: zhaosheng@nuaa.edu.cn
# Date: 2021/09/25
# Description: This script is used to test the model on the test set.
######################################################################################
set -e
. ./path.sh || exit 1
. utils/parse_options.sh || exit 1

# User input password without echo
read -s -p "Enter Password: " password
# echo $password

overwrite=false
master_port=45688
seeds="123 456"
encode_server_type="dguard" # "dguard": means using dguard to encode; "websocket": means using websocket to encode
websocket_nj=40
start_from=0
cpu_nj=64
stage=0
stop_stage=1
mean_embedding=false

# Parse arguments from command line
while getopts ":c:" opt; do
  case $opt in
    c) config_file=$OPTARG;;
    \?) echo "Invalid option -$OPTARG" >&2; exit 1;;
  esac
done

# Check if config file is provided
if [ -z "$config_file" ]; then
  echo "Please provide a config file using the -c option."
  exit 1
fi

# Load all parameters in <config_file>
source $config_file

# if exp_name is not set, use date and time to get a name
if [ -z "$exp_name" ]; then
    todaystr=$(date "+%Y%m%d")
    timestr=$(date "+%H%M%S")
    echo "Please input the EXP name ($todaystr_timestr):"
    read exp_name
    exp_name=${exp_name:-$todaystr_timestr}
fi

# if models is not set, use input to get models
if [ -z "$models" ]; then
    echo "Input models (space-separated):"
    read models
    echo "Models: $models"
fi

# if datasets is not set, use input to get datasets
if [ -z "$datasets" ]; then
    echo "Input datasets (space-separated):"
    read datasets
    echo "Datasets: $datasets"
fi

# If lengths is not set, use input to get lengths
if [ -z "$lengths" ]; then
    echo "Input lengths (space-separated):"
    read lengths
    echo "Lengths: $lengths"
fi

# If GPUs are not set, get GPUs from user
if [ -z "$gpus" ]; then
    # GET available gpus
    all_gpu_num=$(nvidia-smi -L | wc -l)
    
    echo "All GPU num: $all_gpu_num"
    echo "Input GPUs (space-separated, available: 0 to $[$all_gpu_num-1]):"
    read gpus
    nj=$(echo $gpus | wc -w)
    echo "GPUs: $gpus"
    echo "NJ: $nj"
fi

echo "Read config file $config_file"
echo "EXP Name: $exp_name"

# if result/$exp_name already exists, ask if continue, input y/n/Y/N/yes/no/YES/NO
if [ -d result/$exp_name ]; then
    echo "result/$exp_name already exists, do you want to continue? [y/n/r]"
    read ynr 
    case $ynr in
        [Yy]* ) echo "continue";;
        [Rr]* ) echo "Remove old folder" && rm -rf result/$exp_name/*;;
        [Nn]* ) exit;;
        * ) echo "Please answer yes or no.";;
    esac
fi

# copy now script to result/$exp_name/run_shell_<timestr>.sh
timestr=$(date "+%Y%m%d_%H%M%S")
mkdir -p result/$exp_name
cp $0 result/$exp_name/run_shell_$timestr.sh
cp $config_file result/$exp_name/config.sh


######################################################################################
# Function to extract embeddings
######################################################################################
extract_embeddings() {
    local seed=$1
    local model=$2
    local dataset=$3
    local length=$4
    local encode_server_type=$5
    
    echo "Model: $model Datasets: $dataset"
    
    trial_path=$(ls ./trials/$dataset*.trials)
    # remove \r \t \n in trial_path
    trial_path=$(echo $trial_path | tr -d '\r\t\n')    
    mkdir -p result/$exp_name/$model/$dataset/$length
    scp_path=./scp/$dataset.scp
    echo "============================="
    echo "* Trial Paths: $trial_path"
    echo "-----------------------------"
    echo "* Scp Path: $scp_path"
    echo "-----------------------------"
    echo "* Seed: $seed"
    echo "-----------------------------"
    echo "* Length: $length"
    echo "-----------------------------"
    echo "* Type: $encode_server_type"
    echo "-----------------------------"
    # encode_server_type="websocket"
    if [ $encode_server_type == "websocket" ]; then
        echo "Encode server type: websocket"
        # dguard encode
        if [ ! -d result/$exp_name/$model/$trial/embeddings ] || $overwrite; then
            # websocket_nj
            for now_index in $(seq 0 $[$websocket_nj-1]); do
                echo "Now index: $now_index"
                python extract_test_emb_websocket.py --seed $seed --length $length --model_name $model --exp_dir result/$exp_name/$model/$dataset/$length \
                    --start_from=$start_from --data $scp_path --total $websocket_nj --index $now_index || echo "Model:$model Trials:$dataset websocket error" &
            done
        fi
    fi
    wait
    
    if [ $encode_server_type == "dguard" ]; then
        if [ ! -d result/$exp_name/$model/$trial/embeddings ] || $overwrite; then
            # mean_embedding
            if $mean_embedding; then
                echo "Mean embedding"
                torchrun --nproc_per_node=$nj --master_port=$master_port extract_test_emb.py --seed $seed --length $length --model_name $model --exp_dir result/$exp_name/$model/$dataset/$length \
                    --start_from=$start_from --data $scp_path --use_gpu --gpu $gpus --mean_embedding --key $password || echo echo "Model:$model Trials:$dataset torchrun error"
            else
                echo "Not mean embedding"
                torchrun --nproc_per_node=$nj --master_port=$master_port extract_test_emb.py --seed $seed --length $length --model_name $model --exp_dir result/$exp_name/$model/$dataset/$length \
                    --start_from=$start_from --data $scp_path --use_gpu --gpu $gpus --key $password || echo echo "Model:$model Trials:$dataset torchrun error"
            fi
        fi
    fi
    # master_port += 1
    master_port=$[$master_port+1]
    wait
    tiny_save_dir=result/$exp_name/$model/$trial/scores/tiny
    mkdir -p $tiny_save_dir
    for now_index in $(seq 0 $[$cpu_nj-1]); do
        # echo "Now index: $now_index"
        dguard_metrics_multi --total=$cpu_nj --tiny_save_dir=$tiny_save_dir --rank=$now_index --enrol_data result/$exp_name/$model/$dataset/$length/embeddings --test_data result/$exp_name/$model/$dataset/$length/embeddings \
            --scores_all="result/$exp_name/seed_${seed}.csv" --exp_id="${model},${dataset},${length}" --scores_dir result/$exp_name/$model/$dataset/$length/scores --trials $trial_path|| echo echo "Model:$model Trials:$dataset compute_score_metrics error" &
    done
    wait
    dguard_metrics_merge --total=1 --tiny_save_dir=$tiny_save_dir --rank=0 --enrol_data result/$exp_name/$model/$dataset/$length/embeddings --test_data result/$exp_name/$model/$dataset/$length/embeddings \
        --scores_all="result/$exp_name/seed_${seed}.csv" --exp_id="${model},${dataset},${length}" --scores_dir result/$exp_name/$model/$dataset/$length/scores --trials $trial_path --min_precision $min_precision --min_recall $min_recall  || echo echo "Model:$model Trials:$dataset compute_score_metrics error" &

    echo echo "Model:$model Trials:$trial Lenght:$length done!!"
    echo "============================="
}

######################################################################################
# Step 1: Extract embeddings
######################################################################################
if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
    for seed in $seeds; do
        for model in $models; do
            for dataset in $datasets; do
                for length in $lengths; do
                    extract_embeddings $seed $model $dataset $length $encode_server_type
                done
            done
        done
    done
fi

######################################################################################
# Step 2: Get mean csv
######################################################################################
if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
    dguard_mean --csv_folder result/$exp_name --prefix "seed_" --output ./result/$exp_name/mean.csv
fi

######################################################################################
# Step 3: Plot
######################################################################################
if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
    # if len(datasets) >= 6 use dguard_plot_radar ,else use dguard_plot
    if [ ${#datasets[@]} -ge 6 ]; then
        dguard_plot --csv_path ./result/$exp_name/mean.csv --png_path ./result/$exp_name/mean_png/mean
    else
        dguard_plot_radar --csv_path ./result/$exp_name/mean.csv --output_folder ./result/$exp_name/radar
    fi
fi
